import numpy as np
import torch.nn as nn
from mmcv.cnn import normal_init

# from .anchor_head import AnchorHead
from .anchor_head_rbbox_dual import AnchorHeadRbboxDual
from ..registry import HEADS
from ..utils import bias_init_with_prob, ConvModule


@HEADS.register_module
class RetinaHeadRbboxDual(AnchorHeadRbboxDual):

    def __init__(self,
                 num_classes,
                 in_channels,
                 stacked_convs=4,
                 octave_base_scale=4,
                 scales_per_octave=3,
                 conv_cfg=None,
                 norm_cfg=None,
                 **kwargs):
        self.stacked_convs = stacked_convs
        self.octave_base_scale = octave_base_scale
        self.scales_per_octave = scales_per_octave
        self.conv_cfg = conv_cfg
        self.norm_cfg = norm_cfg
        octave_scales = np.array(
            [2**(i / scales_per_octave) for i in range(scales_per_octave)])
        anchor_scales = octave_scales * octave_base_scale
        super(RetinaHeadRbboxDual, self).__init__(
            num_classes, in_channels, anchor_scales=anchor_scales, **kwargs)

    def _init_layers(self):
        self.relu = nn.ReLU(inplace=True)
        self.cls_convs = nn.ModuleList()
        self.reg_convs_h = nn.ModuleList()      # 水平系统
        self.reg_convs_v = nn.ModuleList()      # 垂直系统
        if self.loss_hv_dual is not None:
            self.cls_angle = nn.ModuleList()
        for i in range(self.stacked_convs):
            chn = self.in_channels if i == 0 else self.feat_channels
            self.cls_convs.append(
                ConvModule(
                    chn,
                    self.feat_channels,
                    3,
                    stride=1,
                    padding=1,
                    conv_cfg=self.conv_cfg,
                    norm_cfg=self.norm_cfg))
            self.reg_convs_h.append(
                ConvModule(
                    chn,
                    self.feat_channels,
                    3,
                    stride=1,
                    padding=1,
                    conv_cfg=self.conv_cfg,
                    norm_cfg=self.norm_cfg))
            self.reg_convs_v.append(
                ConvModule(
                    chn,
                    self.feat_channels,
                    3,
                    stride=1,
                    padding=1,
                    conv_cfg=self.conv_cfg,
                    norm_cfg=self.norm_cfg))
            if self.loss_hv_dual is not None:
                self.cls_angle.append(
                    ConvModule(
                        chn,
                        self.feat_channels,
                        3,
                        stride=1,
                        padding=1,
                        conv_cfg=self.conv_cfg,
                        norm_cfg=self.norm_cfg)
                )
        self.retina_cls = nn.Conv2d(
            self.feat_channels,
            self.num_anchors * self.cls_out_channels,
            3,
            padding=1)
        self.retina_reg_h = nn.Conv2d(
            self.feat_channels, self.num_anchors * 5, 3, padding=1)
        self.retina_reg_v = nn.Conv2d(
            self.feat_channels, self.num_anchors * 5, 3, padding=1)
        if self.loss_hv_dual is not None:
            self.retina_cls_angle = nn.Conv2d(
                self.feat_channels,
                self.num_anchors * 2,
                3,
                padding=1)

    def init_weights(self):
        for m in self.cls_convs:
            normal_init(m.conv, std=0.01)
        for m in self.reg_convs_h:
            normal_init(m.conv, std=0.01)
        for m in self.reg_convs_v:
            normal_init(m.conv, std=0.01)
        bias_cls = bias_init_with_prob(0.01)
        normal_init(self.retina_cls, std=0.01, bias=bias_cls)
        if self.loss_hv_dual is not None:
            normal_init(self.retina_cls_angle, std=0.01, bias=bias_cls)
        normal_init(self.retina_reg_h, std=0.01)
        normal_init(self.retina_reg_v, std=0.01)

    def forward_single(self, x):
        cls_feat = x
        reg_feat_h = x
        reg_feat_v = x
        cls_angle_feat = x
        for cls_conv in self.cls_convs:
            cls_feat = cls_conv(cls_feat)
        for reg_conv_h in self.reg_convs_h:
            reg_feat_h = reg_conv_h(reg_feat_h)
        for reg_conv_v in self.reg_convs_v:
            reg_feat_v = reg_conv_v(reg_feat_v)

        cls_score = self.retina_cls(cls_feat)
        bbox_pred_h = self.retina_reg_h(reg_feat_h)
        bbox_pred_v = self.retina_reg_v(reg_feat_v)
        if self.loss_hv_dual is not None:
            for cls_angle_conv in self.cls_angle:
                cls_angle_feat = cls_angle_conv(cls_angle_feat)
            bbox_cls_angle_score = self.retina_cls_angle(cls_angle_feat)
            return cls_score, bbox_pred_h, bbox_pred_v, bbox_cls_angle_score
        return cls_score, bbox_pred_h, bbox_pred_v, None
